'''
Created on Oct 20, 2009

@author: mkiyer
'''
import sys
import operator

def parse_cuffcompare_tmap(fhd):
    '''
    ref_gene_id     ref_id  class_code      cuff_gene_id    cuff_id FMI     RPKM    <unknown>    <unknown>    cov     len     major_iso_id
    ['OTTHUMT00000002844', 'OTTHUMT00000002844', 'c', '3064YAAXX_2.1', '3064YAAXX_2.1.0', '100', '0.164120', '0.000000', '0.000000', '0.108248', '76', '3064YAAXX_2.1.0']
    '''
    # first line is header
    fhd.readline()
    for thisline in fhd:        
        if not thisline:
            continue
        if thisline.startswith('#'):
            continue
        thisline = thisline.strip()
        if not thisline:
            continue
        thisfields = thisline.split('\t')
        t = type('TMapLine', (object,), dict())
        t.line = thisline
        t.ref_gene_id = thisfields[0] if thisfields[0] is not '-' else None
        t.ref_id = thisfields[1] if thisfields[1] is not '-' else None
        t.class_code = thisfields[2]
        t.cuff_gene_id = thisfields[3]
        t.cuff_id = thisfields[4]
        t.fmi = int(thisfields[5])
        t.rpkm = float(thisfields[6])
        t.cov = float(thisfields[9])
        t.length = int(thisfields[10])
        t.major_iso_id = thisfields[11]
        yield t

def histogram(data, xlabel='x', ylabel='y'):
    import matplotlib.pyplot as plt
    import matplotlib.mlab as mlab
    import numpy as np

    data = np.array(data)
    print 'average', np.average(data)

    fig = plt.figure()
    ax = fig.add_subplot(111)
    nbins=250

    # the histogram of the data
    n, bins, patches = ax.hist(data, bins=nbins, facecolor='green', alpha=0.5, log=True, normed=0)
    plt.axvline(x=np.average(data), linewidth=4, color='r')
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    #n, bins, patches = ax.hist(x2, bins=bins, facecolor='red', alpha=0.25)
    #l, = plt.plot(bins, mlab.normpdf(bins, 0.0, 1.0), 'r--', label='fit', linewidth=3)
    #legend([l, patches[0]], ['fit', 'hist'])
    ax.grid(True)
    plt.show()
    fig.close()
    return

def plot_scatter(data, 
                 xlabel='Control Group',
                 ylabel='Experimental Group',
                 title='Coverage scatter plot'):
    import numpy as np
    import matplotlib.pyplot as plt
    import matplotlib.cm as cm
    import math
    
    fig = plt.figure()
    ax = fig.add_subplot(111)
    plt.title(title)

    #sorted_data = sorted(data, operator.itemgetter(0))
    
    xdata = [d[0] for d in data]
    ydata = [d[1] for d in data]
    # size?
    # s = 20*np.log(np.array([x.ratio for x in intervals]))        
    # color?
    # c = [x.category * cm.jet.N / len(category_names)] * len(intervals)
    # c = x.category * cm.jet.N / len(category_names)
    plt.scatter(xdata, ydata, 
                s=30, c='b', marker='o', cmap=None, norm=None,
                vmin=None, vmax=None, alpha=1.0, linewidths=None,
                verts=None)
    #ax.set_xscale('log')
    #ax.set_yscale('log')

    plt.axvline(x=np.average(xdata), linewidth=4, color='r')
    plt.axhline(y=np.average(ydata), linewidth=4, color='g')
    
    print 'x', np.average(xdata)
    print 'y', 'avg', np.average(ydata), 'min', np.min(ydata), 'max', np.max(ydata)
    
    # label=category_names[category])
    #plt.legend()
    # best fit line
    sorted_data = sorted(data, key=operator.itemgetter(0))
    x, y = [d[0] for d in sorted_data], [d[1] for d in sorted_data]    
    m = np.polyfit(x, y, 1)
    yfit = np.polyval(m, x)
    plt.plot(x, yfit, 'g')
    # correlation coefficient
    c = np.corrcoef(xdata, ydata)[0,1]
    r2 = c**2
    plt.figtext(0.815, 0.013, ' r^2=%.3f' % r2, color='black', weight='roman',
               size='small')    
    plt.grid()
    themax = max(x[-1], y[-1])
    plt.axis([-10, themax, -10, themax])
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.show()
    plt.close()

def profile_unknown_transcripts(tmap_file):    
    lengths = set([])
    rpkms = set([])
    rpkm_length = set([])
    for t in parse_cuffcompare_tmap(open(tmap_file)):
        if t.class_code == 'u':
            #lengths.add(t.length)
            #rpkms.add(t.rpkm)
            if t.length >= 300:
                rpkm_length.add((t.length, t.rpkm))
    #histogram(list(lengths))
    #histogram(list(rpkms))
    #plot_scatter(rpkm_length)

def filter_novel_candidates(tmap_file):
    for t in parse_cuffcompare_tmap(open(tmap_file)):
        if t.class_code == 'u' and t.length >= 300:
            print t.line

if __name__ == '__main__':
    #profile_unknown_transcripts(sys.argv[1])
    filter_novel_candidates(sys.argv[1])
    
    